#!/usr/bin/env ruby

require 'fileutils'
require 'open3'
require 'securerandom'
require 'yaml'

DEBUG = false

# Bomb out if we're root
if !Gem.win_platform? && Process.uid.zero?
  puts "Please run #{File.basename($0)} as a non-root user"
  abort
end

@base = '<%= install_dir %>'
@embedded = "#{@base}/embedded"
@bin = "#{@embedded}/bin"
@prev_bin = "#{@embedded}/postgresql-prev/bin"
@framework = "#{@embedded}/framework"

@localconf = "#{ENV['HOME']}/.msf4"
@db = "#{@localconf}/db"
@dbconf = "#{@localconf}/database.yml"
@dbport = 5433
@pg_ctl = "#{@bin}/pg_ctl"

@msf_user = 'msf'
@msftest_user = 'msftest'


def run_cmd(cmd, input = nil)
  exitstatus = 0

  err = out = ""
  Open3.popen3(cmd) do |stdin, stdout, stderr, wait_thr|
    stdin.puts(input) if input
    if DEBUG
      err = stderr.read
      out = stdout.read
    end
    exitstatus = wait_thr.value.exitstatus
  end

  if exitstatus != 0
    if DEBUG
      puts "'#{cmd}' returned #{exitstatus}"
      puts out
      puts err
    end
  end

  exitstatus
end

def run_psql(cmd)
  run_cmd("#{@bin}/psql -p #{@dbport} -c \"#{cmd};\" postgres")
  if DEBUG
    $stderr.puts "#{@bin}/psql -p #{@dbport} -c \"#{cmd};\" postgres"
  end
end

def pw_gen
  SecureRandom.base64(32)
end

def tail(file)
  File.readlines(file).last.to_s.strip
end

def status_db
  if Dir.exist?(@db)
    if run_cmd("#{@pg_ctl} -o \"-p #{@dbport}\" -D #{@db} status") == 0
      puts "Database started at #{@db}"
    else
      puts "Database is not running at #{@db}"
    end
  else
    puts "No database found at #{@db}"
  end
end

def started_db
  if run_cmd("#{@pg_ctl} -o \"-p #{@dbport}\" -D #{@db} status") == 0
    puts "Database already started at #{@db}"
    return true
  end

  print "Starting database at #{@db}..."
  run_cmd("#{@pg_ctl} -o \"-p #{@dbport}\" -D #{@db} -l #{@db}/log start")
  sleep(2)
  if run_cmd("#{@pg_ctl} -o \"-p #{@dbport}\" -D #{@db} status") != 0
    puts 'failed'
    false
  else
    puts 'success'
    true
  end
end

def start_db
  if !Dir.exist?(@db)
    puts "No database found at #{@db}, not starting"
    return
  end

  if File.exist?(@dbconf)
    config = YAML.load(File.read(@dbconf))
    if config["production"] && config["production"]["port"]
      port = config["production"]["port"]
      if port != @dbport
        puts "Using database port #{port}"
        @dbport = port
      end
    end
  end

  while !started_db
    last_log = tail("#{@db}/log")
    puts last_log
    fixed = false
    if last_log =~ /not compatible/
      if ask_yn('Attempt to upgrade the database?')
        fixed = upgrade_db
      end
    end
    if !fixed
      if ask_yn('If your database is corrupt, would you to reinitialize it?')
        fixed = reinit_db
      end
    end
    if !fixed
      if !ask_yn('Database not started, try again?')
        return
      end
    end
  end
end

def stop_db
  if run_cmd("#{@pg_ctl} -o \"-p #{@dbport}\" -D #{@db} status") == 0
    puts "Stopping database at #{@db}"
    run_cmd("#{@pg_ctl} -o \"-p #{@dbport}\" -D #{@db} stop")
  else
    puts "Database is no longer running at #{@db}"
  end
end

def create_db
  puts "Creating database at #{@db}"
  Dir.mkdir(@db)
  run_cmd("#{@bin}/initdb --auth-host=trust --auth-local=trust -E UTF8 #{@db}")

  File.open("#{@db}/pg_hba.conf", 'w') do |f|
    f.puts "host    \"msf\"      \"#{@msf_user}\"      127.0.0.1/32           md5"
    f.puts "host    \"msftest\"  \"#{@msftest_user}\"  127.0.0.1/32           md5"
    f.puts "host    \"postgres\"  \"#{@msftest_user}\"  127.0.0.1/32          trust"
    f.puts "host    \"template1\"   all                127.0.0.1/32           trust"
    if Gem.win_platform?
      f.puts "host    all             all                127.0.0.1/32           trust"
      f.puts "host    all             all                ::1/128                trust"
    else
      f.puts "local   all             all                                       trust"
    end
  end

  File.open("#{@db}/postgresql.conf", 'a') do |f|
    f.puts "port = #{@dbport}"
  end
end

def upgrade_db
  unless File.exist?(@prev_bin)
    puts "An old version of postgresql is not available, run 'msfdb reinit'"
    return false
  end
  prev_db = "#{@db}.prev"
  FileUtils.rm_rf(prev_db)
  FileUtils.mv(@db, prev_db)
  puts "Moved old database to #{prev_db}"
  print "Upgrading..."
  create_db
  if run_cmd("#{@bin}/pg_upgrade -d #{prev_db} -D #{@db} -b #{@prev_bin} -B #{@bin}") == 0
    puts "successful, starting..."
    true
  else
    puts "unsuccessful"
    false
  end
end

def init_db
  if Dir.exist?(@db)
    puts "Found a database at #{@db}, checking to see if it is started"
    start_db
    return
  end

  if File.exist?(@dbconf)
    if !ask_yn("Found database config at #{@dbconf}, do you want to overwrite it?")
      return
    end
  end

  # Generate new database passwords
  msf_pass = pw_gen
  msftest_pass = pw_gen

  # Write a default database config file
  Dir.mkdir(@localconf) unless File.directory?(@localconf)
  File.open(@dbconf, 'w') do |f|
    f.puts <<EOF
development: &pgsql
  adapter: postgresql
  database: msf
  username: #{@msf_user}
  password: #{msf_pass}
  host: 127.0.0.1
  port: #{@dbport}
  pool: 200
  timeout: 5

production: &production
  <<: *pgsql

test:
  <<: *pgsql
  database: msftest
  username: #{@msftest_user}
  password: #{msftest_pass}
EOF
  end

  File.chmod(0640, @dbconf)

  create_db

  start_db

  puts 'Creating database users'
  run_psql("create user #{@msf_user} with password '#{msf_pass}'")
  run_psql("create user #{@msftest_user} with password '#{msftest_pass}'")
  run_psql("alter role msf createdb")
  run_psql("alter role msftest createdb")
  run_psql("alter role #{@msf_user} with password '#{msf_pass}'")
  run_psql("alter role #{@msftest_user} with password '#{msftest_pass}'")
  run_cmd("#{@bin}/createdb -p #{@dbport} -O #{@msf_user} -h 127.0.0.1 -U msf -E UTF-8 -T template0 msf",
    "#{msf_pass}\n#{msf_pass}\n")
  run_cmd("#{@bin}/createdb -p #{@dbport} -O #{@msftest_user} -h 127.0.0.1 -U msftest -E UTF-8 -T template0 msftest",
    "#{msftest_pass}\n#{msftest_pass}\n")

  puts 'Creating initial database schema'
  Dir.chdir(@framework) do
    run_cmd("#{@bin}/bundle exec #{@bin}/rake db:migrate")
  end
end

def ask_yn(question)
  loop do
    print "#{question}: "
    yn = STDIN.gets
    case yn
    when /^[Yy]/
      return true
    when /^[Nn]/
      return false
    else
      puts 'Please answer yes or no.'
    end
  end
end


def delete_db
  if Dir.exist?(@db)
    puts "Deleting all data at #{@db}"
    stop_db
    FileUtils.rm_rf(@db)
    if File.exist?(@dbconf) && ask_yn("Delete database configuration at #{@dbconf}?")
      File.delete(@dbconf)
    end
  else
    puts "No data at #{@db}, doing nothing"
  end
end

def reinit_db
  delete_db
  init_db
end

def usage
  prog = File.basename($0)
  puts <<EOF

Manage a metasploit framework database

  #{prog} init    # initialize the database
  #{prog} reinit  # delete and reinitialize the database
  #{prog} delete  # delete database and stop using it
  #{prog} status  # check database status
  #{prog} start   # start the database
  #{prog} stop    # stop the database
  #{prog} restart # restart the database

EOF
  exit
end

if ARGV.length != 1
  usage
end

case ARGV[0]
when 'init'
  init_db
when 'reinit'
  reinit_db
when 'delete'
  delete_db
when 'status'
  status_db
when 'start'
  start_db
when 'stop'
  stop_db
when 'restart'
  stop_db
  start_db
else
  puts "Error: unrecognized action '#{ARGV[0]}'"
  usage
end
